package com.khalil.kkadmin.ai.rag.knowledgeset.service.impl;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.TypeReference;
import com.khalil.kkadmin.ai.rag.knowledgeset.entity.RagKnowledgeDoc;
import com.khalil.kkadmin.ai.rag.knowledgeset.entity.RagKnowledgeSegment;
import com.khalil.kkadmin.ai.rag.knowledgeset.entity.RagKnowledgeText;
import com.khalil.kkadmin.ai.rag.knowledgeset.service.RagKnowledgeDelegateService;
import com.khalil.kkadmin.ai.rag.knowledgeset.service.RagKnowledgeDocService;
import com.khalil.kkadmin.ai.rag.knowledgeset.service.RagKnowledgeSegmentService;
import com.khalil.kkadmin.ai.rag.knowledgeset.service.RagKnowledgeTextService;
import com.khalil.kkadmin.system.dms.DmsDocumentService;
import com.khalil.kkadmin.system.dms.FileInfo;
import org.apache.commons.io.FilenameUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.InputStreamResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Service
public class RagKnowledgeDelegateServiceImpl implements RagKnowledgeDelegateService {
    @Autowired
    private RagKnowledgeDocService ragKnowledgeDocService;
    @Autowired
    private RagKnowledgeTextService ragKnowledgeTextService;
    @Autowired
    private RagKnowledgeSegmentService ragKnowledgeSegmentService;

    @Autowired
    private DmsDocumentService dmsDocumentService;

    @Override
    public void documentETL(Long knowledgeDocId) {
        RagKnowledgeDoc knowledgeDoc = ragKnowledgeDocService.findById(knowledgeDocId);
        if (knowledgeDoc == null) {
            return;
        }
        Long dmsDocumentId = knowledgeDoc.getDocumentId();
        FileInfo fileInfo = dmsDocumentService.download(dmsDocumentId);

        List<Document> readDocuments = readPdf(fileInfo.getInputStream());
        ragKnowledgeTextService.deleteByRagDocumentId(knowledgeDocId);
        for (int i = 0; i < readDocuments.size(); i++) {
            Document readDocument = readDocuments.get(i);
            String text = readDocument.getText();
            if (text == null || text.isEmpty()) {
                continue;
            }
            RagKnowledgeText ragKnowledgeText = new RagKnowledgeText();
            ragKnowledgeText.setKnowledgeSetId(knowledgeDoc.getKnowledgeSetId());
            ragKnowledgeText.setRagDocumentId(knowledgeDocId);
            String originalFilename = fileInfo.getOriginalFilename();
            String baseName = FilenameUtils.getBaseName(originalFilename);
            ragKnowledgeText.setTitle(baseName + " - " + (i + 1));
            ragKnowledgeText.setSource(RagKnowledgeText.SourceEnum.FILE);
            // 去除多余的空格换行等
            text = text.replaceAll("\\s+", " ");
            text = text.trim();
            ragKnowledgeText.setText(text);
            Map<String, Object> metadata = readDocument.getMetadata();
            metadata.put("fileName", originalFilename);
            ragKnowledgeText.setMetadata(JSON.toJSONString(metadata));
            ragKnowledgeTextService.create(ragKnowledgeText);

            // 分段
            List<RagKnowledgeSegment> knowledgeSegments = splitKnowledgeSegmentList(ragKnowledgeText);
        }
        knowledgeDoc.setTextExtractStatus(RagKnowledgeDoc.TEXT_EXTRACT_STATUS_COMPLETE);
        ragKnowledgeDocService.update(knowledgeDoc);
    }

    @Override
    public List<RagKnowledgeSegment> splitKnowledgeTextIdList(List<Long> knowledgeTextIdList) {
        List<RagKnowledgeSegment> ragKnowledgeSegmentList = new ArrayList<>();
        for (Long knowledgeTextId : knowledgeTextIdList) {
            RagKnowledgeText knowledgeText = ragKnowledgeTextService.findById(knowledgeTextId);
            List<RagKnowledgeSegment> segmentList = splitKnowledgeSegmentList(knowledgeText);
            ragKnowledgeSegmentList.addAll(segmentList);
        }
        return ragKnowledgeSegmentList;
    }

    @Transactional
    public List<RagKnowledgeSegment> splitKnowledgeSegmentList(RagKnowledgeText knowledgeText) {
        List<Document> documentList = splitDocumentListFromKnowledgeTextList(List.of(knowledgeText));

        List<RagKnowledgeSegment> ragKnowledgeSegmentList = new ArrayList<>();
        ragKnowledgeSegmentService.deleteByKnowledgeTextId(knowledgeText.getId());
        for (Document document : documentList) {
            RagKnowledgeSegment knowledgeSegment = new RagKnowledgeSegment();
            knowledgeSegment.setKnowledgeSetId(knowledgeText.getKnowledgeSetId());
            knowledgeSegment.setKnowledgeTextId(knowledgeText.getId());
            knowledgeSegment.setText(document.getText());
            knowledgeSegment.setMetadata(JSON.toJSONString(document.getMetadata()));
            ragKnowledgeSegmentService.create(knowledgeSegment);
            ragKnowledgeSegmentList.add(knowledgeSegment);
        }
        return ragKnowledgeSegmentList;
    }

    public List<Document> splitDocumentListFromKnowledgeText(Long knowledgeTextId) {
        RagKnowledgeText knowledgeText = ragKnowledgeTextService.findById(knowledgeTextId);
        if (knowledgeText == null) {
            return Collections.emptyList();
        }
        return splitDocumentListFromKnowledgeTextList(List.of(knowledgeText));
    }

    private List<Document> splitDocumentListFromKnowledgeTextList(List<RagKnowledgeText> knowledgeTextList) {
        List<Document> documentList = knowledgeTextList.stream().map(knowledgeText -> {
            return Document.builder()
                    .text(knowledgeText.getText())
                    .metadata(JSON.parseObject(knowledgeText.getMetadata(), new TypeReference<Map<String, Object>>() {
                    }))
                    .build();
        }).collect(Collectors.toList());
        return splitDocuments(documentList);
    }

    private List<Document> splitDocuments(List<Document> documents) {
        TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();
        return tokenTextSplitter.split(documents);
    }

    private List<Document> readPdf(InputStream inputStream) {
        Resource resource = new InputStreamResource(inputStream);
        return readPdf(resource);
    }

    private List<Document> readPdf(Resource resource) {
        PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(resource,
                PdfDocumentReaderConfig.builder()
                        .withPageTopMargin(0)
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.builder()
                                .withNumberOfTopTextLinesToDelete(0)
                                .build())
                        .withPagesPerDocument(1)
                        .build());

        return pdfReader.read();
    }
}
